# -*- coding: utf-8 -*-
import time
import cv2
import json
import os
import faceppApi
import landmarksTrans
import onMouse

winName = "label window"
labeledFilesInTxt = []

def draw(filepath, outputimgfolder, facebox, ldslist):

    img = cv2.imread(filepath)
    bak = img.copy()

    p0 = (facebox[0], facebox[1])
    p1 = (facebox[2], facebox[3])
    cv2.rectangle(img, p0, p1, (0, 255, 0), 3)

    for i in range(len(ldslist)):
        p = (int(ldslist[i][0]), int(ldslist[i][1]))
        cv2.circle(img, p, 2, (0,0,255), 1)
        #cv2.putText(img, str(i), p, cv2.FONT_HERSHEY_SIMPLEX, 0.35, (0,255,0), 1)

    # backup tmp img
    filename = os.path.basename(filepath)
    outpath = os.path.join(outputimgfolder, "tmp", filename)
    cv2.imwrite(outpath, img)

    # show face roi
    # if landmarks not ok, you can relabel them
    roi = bak[facebox[1]:facebox[3], facebox[0]:facebox[2]]
    w = facebox[2] - facebox[0]
    h = facebox[3] - facebox[1]

    res = cv2.resize(roi, (720, 720))
    for i in range(len(ldslist)):
        p = (int((ldslist[i][0] - facebox[0]) * 720 / w), int((ldslist[i][1] - facebox[1]) * 720 / h))
        cv2.circle(res, p, 2, (0,0,255), 1)
    
    xylist = onMouse.run(res)
    for i in range(len(xylist)):
        x = int(xylist[i][0] * w / 720 + facebox[0])
        y = int(xylist[i][1] * h / 720 + facebox[1])
        xylist[i] = (x, y)
    #print("onMouse return xylist: ", xylist)

    return xylist


def getFaceBox(jsonstr):
    facebox = []

    resultdict = json.loads(jsonstr)
    if len(resultdict["faces"]) > 0:
        #print("sssssssdddddddddddddddddddddddddddddddd,", resultdict)
        x0 = int(resultdict["faces"][0]["face_rectangle"]["left"])
        y0 = int(resultdict["faces"][0]["face_rectangle"]["top"])
        w = int(resultdict["faces"][0]["face_rectangle"]["width"])
        h = int(resultdict["faces"][0]["face_rectangle"]["height"])

        x1 = x0 + w
        y1 = y0 + h
        facebox = [x0, y0, x1, y1]

    return facebox


def getHeadPose(jsonstr):
    headpose = []

    resultdict = json.loads(jsonstr)
    #print("sssssssdddddddddddddddddddddddddddddddd,", resultdict)
    yawangle = float(resultdict["faces"][0]["attributes"]["headpose"]["yaw_angle"])
    pitchangle = float(resultdict["faces"][0]["attributes"]["headpose"]["pitch_angle"])
    rollangle = float(resultdict["faces"][0]["attributes"]["headpose"]["roll_angle"])
    headpose = [yawangle, pitchangle, rollangle]

    return headpose


def saveImg(filepath, facebox, outputimgfolder):
    """
    save ROI into file
    """
    filename = os.path.basename(filepath)

    img = cv2.imread(filepath)
    x0 = facebox[0]
    y0 = facebox[1]
    x1 = facebox[2]
    y1 = facebox[3]
    roi = img[y0:y1, x0:x1]
    #print("llllllllllllllllllllllllllllll roi: ", type(roi), roi.shape)

    # saved face imgs are all 112x112
    dst = cv2.resize(roi, (112, 112))

    #print("llllllllllllllllllllllllllllll dst: ", type(dst), dst.shape)
    outpath = os.path.join(outputimgfolder, "imgs", filename)
    print("outpath: ", outpath)
    cv2.imwrite(outpath, dst)


def saveLabelInfo(filepath, outputimgfolder, facebox, headpose, lds98list):
    """
    save label info into txt file
    """
    x0 = facebox[0]
    y0 = facebox[1]
    x1 = facebox[2]
    y1 = facebox[3]
    w = x1 - x0
    h = y1 - y0

    filename = os.path.basename(filepath)
    fullstring = "imgs/" + filename

    for i in range(len(lds98list)):
        x = lds98list[i][0]
        y = lds98list[i][1]
        
        if x == -256.0:
            xo = -256.0
            yo = -256.0
        else:
            xo = (x - x0) / w
            yo = (y - y0) / h
       
        fullstring = fullstring + " " + str(xo)
        fullstring = fullstring + " " + str(yo)
    
    fullstring = fullstring + " " + "0 0 0 0 0 0"
    fullstring = fullstring + " " + str(headpose[0]) 
    fullstring = fullstring + " " + str(headpose[1]) 
    fullstring = fullstring + " " + str(headpose[2]) 

    # important!!! 
    # mode 0: pose+68lds
    # mode -1: only pose
    # mode -2: 98lds
    fullstring = fullstring + " " + "0" 
    fullstring = fullstring + "\n" 
    #print("fullstring: ", fullstring)

    # save into list.txt file
    txtfile = os.path.join(outputimgfolder, "list.txt")
    with open(txtfile, "a") as f:
        f.write(fullstring)
    f.close()
   

def init(outputimgfolder):
    """
    mkdirs
    """
    # save face-box-imgs
    fold = os.path.join(outputimgfolder, "imgs")
    isExists = os.path.exists(fold)
    if not isExists:
        os.makedirs(fold)
        time.sleep(0.5)

    # save tmp-imgs used for check
    fold = os.path.join(outputimgfolder, "tmp")
    isExists = os.path.exists(fold)
    if not isExists:
        os.makedirs(fold)
        time.sleep(0.5)

    # read list.txt
    txtfile = os.path.join(outputimgfolder, "list.txt")
    isExists = os.path.exists(txtfile)
    if not isExists:
        return

    with open(txtfile, "r") as f:
        line = f.readline()
        content = (line.split("-256.0"))[0]
        labeledFilesInTxt.append(content)
        while line:
            line = f.readline()
            content = (line.split("-256.0"))[0]
            labeledFilesInTxt.append(content)
    f.close()


def run(inimgfolder, outputimgfolder):
    """
    main process
    标注内容：人脸box，人脸landmarks, head pose 
    """
    init(outputimgfolder)

    for root, dirs, files in os.walk(inimgfolder):
        for f in files:

            # check if labeled
            startlabel = True 
            tmpstring = "imgs/" + f
            for i in range(len(labeledFilesInTxt)):
                if tmpstring in labeledFilesInTxt[i]:
                    print(tmpstring + " is labeled.")
                    labeledFilesInTxt.pop(i)
                    startlabel = False
                    break
            if not startlabel:
                continue

            # get result from facePlusPlus Api
            filepath = os.path.join(root, f)
            print("processing ", filepath)

            jsonstr = faceppApi.faceppApi(filepath)
            #print("xxxxxxxxxxxxxxxxx:", jsonstr)

            if len(jsonstr) > 0:
                # get face box 
                facebox = getFaceBox(jsonstr)
                print("facebox: ", facebox)

                if len(facebox) == 0:
                    print("warning!!! not found face!!! img: ", filepath)
                    continue

                # get head pose
                headpose = getHeadPose(jsonstr)
                print("headpose: ", headpose)

                # get face landmarks
                lds83list = landmarksTrans.getlandmarks(jsonstr)

                lds68list = landmarksTrans.transformLds83to68(lds83list)
                lds98list = landmarksTrans.transformLds68to98(lds68list)
                #print(lds98list)
                #print("lds98list len: ", len(lds98list))
                modifiedLds = draw(filepath, outputimgfolder, facebox, lds98list)

                # save face img
                saveImg(filepath, facebox, outputimgfolder)

                # save label info
                landmarksTrans.modifyEyeLds(lds98list, modifiedLds)
                saveLabelInfo(filepath, outputimgfolder, facebox, headpose, lds98list)
            else:
                print("warning: no api return...")


if __name__ == "__main__":
    """
    调用face++的api接口，来实现自动标注, 用于PFPLD模型的训练
    标注内容：人脸box，人脸landmarks, head pose 
    2022.4.16
    """
    inimgfolder = "/mnt/sdb/datasets_myself/20220418_PFPLD_Ilabel/originalPhotos" #"/home/aiot/Desktop/test_in"
    outimgfolder = "/mnt/sdb/datasets_myself/20220418_PFPLD_Ilabel/dataset" #"/home/aiot/Desktop/out"

    run(inimgfolder, outimgfolder)

